##########################################################################################

library(data.table)
library(optparse)
library(ggplot2)
library(RColorBrewer)
library(ggsci)
library("scales")
library(dplyr)

##########################################################################################

option_list <- list(
    make_option(c("--input_file"), type = "character"),
    make_option(c("--smg_file"), type = "character"),
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    input_file <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/smgs/compute_dn_ds.csv"
    smg_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutsig_check/smg.list"
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/finalPlot/revise/smgs"
}

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_file <- opt$input_file
smg_file <- opt$smg_file
out_path <- opt$out_path

##########################################################################################

dat_mut <- data.frame(fread(input_file))
dat_smg <- data.frame(fread(smg_file))

##########################################################################################

tmp <- dat_mut[grep( "CDKN2A" , dat_mut$gene_name ),]
tmp <- tmp %>%
group_by(Class) %>%
summarize( n_syn = sum(n_syn) , n_mis = sum(n_mis) , n_non = sum(n_non) , n_spl = sum(n_spl) , n_ind = sum(n_ind) ,
     wmis_cv = max(wmis_cv) , wnon_cv = max(wnon_cv) , wspl_cv = max(wspl_cv) , wind_cv = max(wind_cv) , 
     pmis_cv = min(pmis_cv) , ptrunc_cv = min(ptrunc_cv) , pallsubs_cv = min(pallsubs_cv) , pind_cv = min(pind_cv) , 
     qmis_cv = min(qmis_cv) , qtrunc_cv = min(qtrunc_cv) , qallsubs_cv = min(qallsubs_cv) , qind_cv = min(qind_cv) , 
     pglobal_cv = min(pglobal_cv) , qglobal_cv = min(qglobal_cv))
tmp$gene_name <- "CDKN2A"
tmp <- data.frame(tmp)
tmp <- tmp[,colnames(dat_mut)]

##########################################################################################
## 提取SMG
dat_use <- subset(dat_mut , gene_name %in% dat_smg$Gene_Symbol )
dat_use <- rbind(dat_use , tmp)

images_name <- paste0(out_path , "/compute_dn_ds.smg.csv")
write.csv(dat_use , images_name , row.names = F , quote = F)

##########################################################################################
## 宽转长
dat_use_mis <- dat_use[,c("gene_name" , "n_mis" , "wmis_cv" , "qmis_cv" , "qglobal_cv" , "Class")]
dat_use_ind <- dat_use[,c("gene_name" , "n_ind" , "wind_cv" , "qind_cv" , "qglobal_cv" , "Class")]
dat_use_non <- dat_use[,c("gene_name" , "n_non" , "wnon_cv" , "qtrunc_cv" , "qglobal_cv" , "Class")]
dat_use_spl <- dat_use[,c("gene_name" , "n_spl" , "wspl_cv" , "qtrunc_cv" , "qglobal_cv" , "Class")]
## splice site和nonsense统称为trunc
dat_use_trunc <- data.frame( 
    gene_name = dat_use_non$gene_name ,
    n_trunc = dat_use_non$n_non + dat_use_spl$n_spl ,
    wtrunc_cv = dat_use_non$wnon_cv + dat_use_spl$wspl_cv ,
    qtrunc_cv = dat_use_non$qtrunc_cv ,
    qglobal_cv = dat_use_non$qglobal_cv ,
    Class = dat_use_non$Class
    )

dat_use_mis$mut_type <- "Missense"
dat_use_trunc$mut_type <- "Nonsense & splice"
dat_use_ind$mut_type <- "Indel"

col_n <- c("gene_name" , "n" , "wcv" , "q" , "qglobal_cv" , "Class" , "mut_type")
colnames(dat_use_mis) <- col_n
colnames(dat_use_trunc) <- col_n
colnames(dat_use_ind) <- col_n

dat_combine <- rbind( dat_use_mis , dat_use_trunc , dat_use_ind )

##########################################################################################
## 每个基因选择最显著的类别
res <- c()
for(gene in unique(dat_combine$gene_name)){
    for(class in unique(dat_combine$Class)){
        tmp <- subset(dat_combine , gene_name == gene & Class == class )
        mutN <- tmp[which.min(tmp$q),"mut_type"]
        tmp <- subset(dat_combine , gene_name == gene & Class == class & mut_type == mutN)
        res <- rbind(res , tmp)
    }       
}

##########################################################################################

gene_order <- c("TP53","ARID1A","CDH1","APC","SMAD4","MUC6","PIK3CA",
    "CTNNB1","RHOA","ERBB2","CFTR","KRAS","MAP2K7","ARID2",
    "RNF43","TGFBR2","BMP6","FBXW7","CDKN2A","MTRR")
gene_order <- gene_order[length(gene_order):1]

## 可视化
result_use <- dat_combine
result_use$percent_pos <- result_use$wcv + 1
result_use$percent_pos <- as.numeric(result_use$percent_pos)
result_use$gene_name <- factor( result_use$gene_name , levels = gene_order , order = T )
result_use$mut_type <- factor( result_use$mut_type , levels = c("Missense" , "Nonsense & splice" , "Indel") , order = T )
result_use$p_text <- ifelse( result_use$q < 0.05 , "*" , "" )

#print(result_use)
result_use$Class <- factor( result_use$Class , levels = c("IM" , "IGC" , "DGC") , order = T )
result_use$wcv <- ifelse( result_use$wcv > 100 , 100 , result_use$wcv )

## 展示不同基因集合
maintained_gene_igc <- c("TP53","APC","PIK3CA")
maintained_gene_dgc <- c("TP53","APC","PIK3CA","CDH1")
im_favored_gene <- c("MUC6" , "CFTR" , "BMP6" , "MTRR")
gc_favored_gene_igc <- c("ARID1A","CDH1","SMAD4","CTNNB1","RHOA","ERBB2","MAP2K7","ARID2","RNF43","TGFBR2","FBXW7")
gc_favored_gene_dgc <- c("ARID1A","SMAD4","CTNNB1","RHOA","ERBB2","MAP2K7","ARID2","RNF43","TGFBR2","FBXW7")

gene_order <- c(maintained_gene_dgc , im_favored_gene , gc_favored_gene_dgc)
gene_order <- c(gene_order , "KRAS" , "CDKN2A")
gene_order <- gene_order[length(gene_order):1]
result_use$gene_name <- factor( result_use$gene_name , levels = gene_order , order = T )

for(classN in unique(result_use$Class)){
    if(1!=1){
        if(classN == "IM"){
            show_gene <- c( maintained_gene_dgc , im_favored_gene )
        }else if(classN == "IGC"){
            show_gene <- maintained_gene_igc
        }else if(classN == "DGC"){
            show_gene <- maintained_gene_dgc
        }
    }
    if(classN == "IM"){
        maintained_gene <- maintained_gene_dgc
        gc_favored_gene <- gc_favored_gene_dgc
    }else if(classN == "IGC"){
        maintained_gene <- maintained_gene_igc
        gc_favored_gene <- gc_favored_gene_igc
    }else if(classN == "DGC"){
        maintained_gene <- maintained_gene_dgc
        gc_favored_gene <- gc_favored_gene_dgc
    }

    show_gene <- unique(subset(result_use , q < 0.05 & Class == classN)$gene_name)
    result_use_tmp <- subset(result_use , gene_name %in% show_gene & Class == classN)
    result_use_tmp$gene_type <- ifelse( result_use_tmp$gene_name %in% im_favored_gene , "IM favored" , "" )
    result_use_tmp$gene_type <- ifelse( result_use_tmp$gene_name %in% maintained_gene , "Maintained" , result_use_tmp$gene_type)
    result_use_tmp$gene_type <- ifelse( result_use_tmp$gene_name %in% gc_favored_gene , "GC favored" , result_use_tmp$gene_type)
    result_use_tmp$gene_type <- factor( result_use_tmp$gene_type , levels = c("Maintained" , "IM favored" , "GC favored") )

    p <- ggplot(result_use_tmp,mapping = aes(x = wcv , y = gene_name , fill = mut_type)) +
    geom_bar(stat = 'identity', position = 'dodge' , width = 0.8 , color = 'black') + 
    scale_x_continuous( limits = c(0 , 100) , breaks = c( 1 , 20 , 50 , 80 , 100 ) , labels = c( 1 , 20 , 50 , 80 , ">100" ) , position="top") +
    theme_bw() +
    #facet_grid(rows = vars(gene_type) , scales = "free" ) +
    #geom_text(aes(label= p_text , y = gene_name , x = wcv ),size= 5 ,family="Helvetica")+
    #geom_text(aes(label=Type , x = 18 , y = label_pos ),size=5,family="Helvetica")+
    ylab("") +
    xlab('dN/dS') +
    scale_fill_aaas() +
    theme(
        title =element_text(size=4, face='bold'),
        legend.title = element_blank(),
        legend.text = element_text(size = 8),
        legend.key.width = unit(1, "cm"),
        legend.key.height = unit(1, "cm"),
        legend.position = "bottom" ,
        axis.text.y = element_text(size = 11 , color = 'black' , family="Helvetica" ) ,
        axis.title.x =  element_text(size = 14 , color = 'black' ) ,
        axis.text.x =  element_text(size = 11 , color = 'black' , family="Helvetica") ,
        axis.ticks.length = unit(0.2, "cm") ,
        panel.grid.major =element_blank(),
        panel.grid=element_blank() ,
        axis.line.x=element_line(linetype=1,color="black",size=0.2)
    )

    width <- 8
    height <- length(show_gene)/1.5 + 1
    images_name <- paste0(out_path , "/compute_dn_ds.",classN,".pdf")
    ggsave( images_name , p , width = width , height = height )

}

##########################################################################################
## 画在一张图
p <- ggplot(result_use,mapping = aes(x = wcv , y = gene_name , fill = Class)) +
geom_bar(stat = 'identity', position = 'dodge' , width = 0.8) + 
scale_x_continuous( limits = c(0 , 110) , breaks = c( 1 , 20 , 50 , 90 , 100 ) , labels = c( 1 , 20 , 50 , 90 , ">100" ) , position="top") +
theme_bw() +
facet_grid(cols = vars(Class) , scales = "free" ) +
geom_text(aes(label= p_text , y = gene_name , x = 110 ),size= 5 ,family="Helvetica")+
ylab("") +
xlab('dN/dS') +
scale_fill_npg() +
theme(
    title =element_text(size=4, face='bold'),
    legend.title = element_blank(),
    legend.text = element_text(size = 8),
    legend.key.width = unit(1, "cm"),
    legend.key.height = unit(1, "cm"),
    legend.position = "none" ,
    axis.text.y = element_text(size = 11 , color = 'black' , family="Helvetica" ) ,
    axis.title.x =  element_text(size = 14 , color = 'black' ) ,
    axis.text.x =  element_text(size = 11 , color = 'black' , family="Helvetica") ,
    axis.ticks.length = unit(0.2, "cm") ,
    panel.grid.major =element_blank(),
    panel.grid=element_blank() ,
    axis.line.x=element_line(linetype=1,color="black",size=0.2)
)

width <- 8
height <- 6
images_name <- paste0(out_path , "/compute_dn_ds.pdf")
ggsave( images_name , p , width = width , height = height )
